/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.taskexecutor;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.configuration.CoreOptions;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.blob.PermanentBlobService;
import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.execution.librarycache.BlobLibraryCacheManager;
import org.apache.flink.runtime.execution.librarycache.LibraryCacheManager;
import org.apache.flink.runtime.io.disk.iomanager.IOManager;
import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync;
import org.apache.flink.runtime.io.network.TaskEventDispatcher;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.rpc.FatalErrorHandler;
import org.apache.flink.runtime.shuffle.ShuffleEnvironment;
import org.apache.flink.runtime.shuffle.ShuffleEnvironmentContext;
import org.apache.flink.runtime.shuffle.ShuffleServiceLoader;
import org.apache.flink.runtime.state.TaskExecutorLocalStateStoresManager;
import org.apache.flink.runtime.taskexecutor.slot.TaskSlotTable;
import org.apache.flink.runtime.taskexecutor.slot.TaskSlotTableImpl;
import org.apache.flink.runtime.taskexecutor.slot.TimerService;
import org.apache.flink.runtime.taskmanager.Task;
import org.apache.flink.runtime.taskmanager.UnresolvedTaskManagerLocation;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.FlinkException;
import org.apache.flink.util.Preconditions;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;

/**
 * Container for {@link TaskExecutor} services such as the {@link MemoryManager}, {@link IOManager},
 * {@link ShuffleEnvironment}. All services are exclusive to a single {@link TaskExecutor}.
 * Consequently, the respective {@link TaskExecutor} is responsible for closing them.
 */
public class TaskManagerServices {
	private static final Logger LOG = LoggerFactory.getLogger(TaskManagerServices.class);

	@VisibleForTesting
	public static final String LOCAL_STATE_SUB_DIRECTORY_ROOT = "localState";

	/**
	 * TaskManager services.
	 */
	private final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation;
	private final long managedMemorySize;
	private final IOManager ioManager;
	private final ShuffleEnvironment<?, ?> shuffleEnvironment;
	private final KvStateService kvStateService;
	private final BroadcastVariableManager broadcastVariableManager;
	private final TaskSlotTable<Task> taskSlotTable;
	private final JobTable jobTable;
	private final JobLeaderService jobLeaderService;
	private final TaskExecutorLocalStateStoresManager taskManagerStateStore;
	private final TaskEventDispatcher taskEventDispatcher;
	private final ExecutorService ioExecutor;
	private final LibraryCacheManager libraryCacheManager;

	TaskManagerServices(UnresolvedTaskManagerLocation unresolvedTaskManagerLocation, long managedMemorySize, IOManager ioManager,
		ShuffleEnvironment<?, ?> shuffleEnvironment, KvStateService kvStateService, BroadcastVariableManager broadcastVariableManager,
		TaskSlotTable<Task> taskSlotTable, JobTable jobTable, JobLeaderService jobLeaderService,
		TaskExecutorLocalStateStoresManager taskManagerStateStore, TaskEventDispatcher taskEventDispatcher, ExecutorService ioExecutor,
		LibraryCacheManager libraryCacheManager) {

		this.unresolvedTaskManagerLocation = Preconditions.checkNotNull(unresolvedTaskManagerLocation);
		this.managedMemorySize = managedMemorySize;
		this.ioManager = Preconditions.checkNotNull(ioManager);
		this.shuffleEnvironment = Preconditions.checkNotNull(shuffleEnvironment);
		this.kvStateService = Preconditions.checkNotNull(kvStateService);
		this.broadcastVariableManager = Preconditions.checkNotNull(broadcastVariableManager);
		this.taskSlotTable = Preconditions.checkNotNull(taskSlotTable);
		this.jobTable = Preconditions.checkNotNull(jobTable);
		this.jobLeaderService = Preconditions.checkNotNull(jobLeaderService);
		this.taskManagerStateStore = Preconditions.checkNotNull(taskManagerStateStore);
		this.taskEventDispatcher = Preconditions.checkNotNull(taskEventDispatcher);
		this.ioExecutor = Preconditions.checkNotNull(ioExecutor);
		this.libraryCacheManager = Preconditions.checkNotNull(libraryCacheManager);
	}

	// --------------------------------------------------------------------------------------------
	//  Getter/Setter
	// --------------------------------------------------------------------------------------------

	long getManagedMemorySize() {
		return managedMemorySize;
	}

	public IOManager getIOManager() {
		return ioManager;
	}

	public ShuffleEnvironment<?, ?> getShuffleEnvironment() {
		return shuffleEnvironment;
	}

	public KvStateService getKvStateService() {
		return kvStateService;
	}

	public UnresolvedTaskManagerLocation getUnresolvedTaskManagerLocation() {
		return unresolvedTaskManagerLocation;
	}

	public BroadcastVariableManager getBroadcastVariableManager() {
		return broadcastVariableManager;
	}

	public TaskSlotTable<Task> getTaskSlotTable() {
		return taskSlotTable;
	}

	public JobTable getJobTable() {
		return jobTable;
	}

	public JobLeaderService getJobLeaderService() {
		return jobLeaderService;
	}

	public TaskExecutorLocalStateStoresManager getTaskManagerStateStore() {
		return taskManagerStateStore;
	}

	public TaskEventDispatcher getTaskEventDispatcher() {
		return taskEventDispatcher;
	}

	public Executor getIOExecutor() {
		return ioExecutor;
	}

	public LibraryCacheManager getLibraryCacheManager() {
		return libraryCacheManager;
	}

	// --------------------------------------------------------------------------------------------
	//  Shut down method
	// --------------------------------------------------------------------------------------------

	/**
	 * Shuts the {@link TaskExecutor} services down.
	 */
	public void shutDown() throws FlinkException {

		Exception exception = null;

		try {
			taskManagerStateStore.shutdown();
		} catch(Exception e) {
			exception = e;
		}

		try {
			ioManager.close();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			shuffleEnvironment.close();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			kvStateService.shutdown();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			taskSlotTable.close();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			jobLeaderService.stop();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			ioExecutor.shutdown();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			jobTable.close();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		try {
			libraryCacheManager.shutdown();
		} catch(Exception e) {
			exception = ExceptionUtils.firstOrSuppressed(e, exception);
		}

		taskEventDispatcher.clearAll();

		if(exception != null) {
			throw new FlinkException("Could not properly shut down the TaskManager services.", exception);
		}
	}

	// --------------------------------------------------------------------------------------------
	//  Static factory methods for task manager services
	// --------------------------------------------------------------------------------------------

	/**
	 * Creates and returns the task manager services.
	 *
	 * @param taskManagerServicesConfiguration task manager configuration
	 * @param permanentBlobService             permanentBlobService used by the services
	 * @param taskManagerMetricGroup           metric group of the task manager
	 * @param ioExecutor                       executor for async IO operations
	 * @param fatalErrorHandler                to handle class loading OOMs
	 * @return task manager components
	 * @throws Exception
	 */
	public static TaskManagerServices fromConfiguration(TaskManagerServicesConfiguration taskManagerServicesConfiguration,
		PermanentBlobService permanentBlobService, MetricGroup taskManagerMetricGroup, ExecutorService ioExecutor,
		FatalErrorHandler fatalErrorHandler) throws Exception {

		// pre-start checks
		checkTempDirs(taskManagerServicesConfiguration.getTmpDirPaths());

		// 1. 初始化 TaskEventDispatcher
		// Flink 运行的是 流式任务:  StreamTask （OperatorChain Pipline   输入给我，调用 Transformation执行，输出结果）
		final TaskEventDispatcher taskEventDispatcher = new TaskEventDispatcher();

		// 2. 初始化 IOManagerASync
		// 启动一定数量的 WriterThread 和 ReaderThread
		// start the I/O manager, it will create some temp directories.
		final IOManager ioManager = new IOManagerAsync(taskManagerServicesConfiguration.getTmpDirPaths());

		// 3. 初始化 NettyShuffleEnvironment
		// 上游 StreamTask 和 下游 StreamTask 有 shuffle 动作。
		// 在这个动作过程中，肯定需要很多的一些组件来为其服务， 现在创建的这个 NettyShuffleEnvironment 对象为将来的 SHuffle 提供各种组件创建的支撑
		// shuffleEnvironment = NettyShuffleEnvironment
		final ShuffleEnvironment<?, ?> shuffleEnvironment = createShuffleEnvironment( // TODO
			taskManagerServicesConfiguration,
			taskEventDispatcher,
			taskManagerMetricGroup,
			ioExecutor);

		// 启动过程中，启动了 Netty 服务端 和 客户端，负责 IO 的
		final int listeningDataPort = shuffleEnvironment.start();

		// 4. 初始化 KVStageService 状态管理服务
		final KvStateService kvStateService = KvStateService.fromConfiguration(taskManagerServicesConfiguration);
		kvStateService.start();

		final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new UnresolvedTaskManagerLocation(
			taskManagerServicesConfiguration.getResourceID(), taskManagerServicesConfiguration.getExternalAddress(),
			// we expose the task manager location with the listening port if the external data port is not explicitly defined
			taskManagerServicesConfiguration.getExternalDataPort() > 0 ?
				taskManagerServicesConfiguration.getExternalDataPort() : listeningDataPort);

		// 5. 初始化 BroadCastVariableManager = BroadcastVariableManager
		final BroadcastVariableManager broadcastVariableManager = new BroadcastVariableManager();

		// 6. 初始化 TaskSlotTable = TaskSlotTableImpl
		// 当前节点能提供很多的 Slot， 也会执行很多的 Task
		// 到底哪些 Task 对应到 哪些 Slot 执行，这里面通过一张表进行管理
		final TaskSlotTable<Task> taskSlotTable = createTaskSlotTable( // TODO
			taskManagerServicesConfiguration.getNumberOfSlots(),
			taskManagerServicesConfiguration.getTaskExecutorResourceSpec(),
			taskManagerServicesConfiguration.getTimerServiceShutdownTimeout(),
			taskManagerServicesConfiguration.getPageSize(),
			ioExecutor);

		// 7. 初始化 JobTable = DefaultJobTable
		final JobTable jobTable = DefaultJobTable.create();

		// 8. 初始化 JobLeaderService（即JobMaster）
		final JobLeaderService jobLeaderService = new DefaultJobLeaderService(
			unresolvedTaskManagerLocation,
			taskManagerServicesConfiguration.getRetryingRegistrationConfiguration()
		);

		// 9. 初始化 StateRootDirectory 和 StateRootDirectoryFile
		final String[] stateRootDirectoryStrings = taskManagerServicesConfiguration.getLocalRecoveryStateRootDirectories();
		final File[] stateRootDirectoryFiles = new File[stateRootDirectoryStrings.length];
		for(int i = 0; i < stateRootDirectoryStrings.length; ++i) {
			stateRootDirectoryFiles[i] = new File(stateRootDirectoryStrings[i], LOCAL_STATE_SUB_DIRECTORY_ROOT);
		}

		// 10. 初始化 TaskExecutorLocalStateStoresManager
		final TaskExecutorLocalStateStoresManager taskStateManager = new TaskExecutorLocalStateStoresManager(
			taskManagerServicesConfiguration.isLocalRecoveryEnabled(), stateRootDirectoryFiles, ioExecutor);

		final boolean failOnJvmMetaspaceOomError = taskManagerServicesConfiguration.getConfiguration()
			.getBoolean(CoreOptions.FAIL_ON_USER_CLASS_LOADING_METASPACE_OOM);

		// 11. 初始化 LibraryCacheManager
		//   第一个参数： PermanentBlobService
		//   第二个参数： DefaultClassLoaderFactory
		final LibraryCacheManager libraryCacheManager = new BlobLibraryCacheManager(
			permanentBlobService,
			BlobLibraryCacheManager.defaultClassLoaderFactory(
				taskManagerServicesConfiguration.getClassLoaderResolveOrder(),
				taskManagerServicesConfiguration.getAlwaysParentFirstLoaderPatterns(),
				failOnJvmMetaspaceOomError ? fatalErrorHandler : null)
		);

		// 返回： TaskManagerServices
		return new TaskManagerServices(
			unresolvedTaskManagerLocation,
			taskManagerServicesConfiguration.getManagedMemorySize().getBytes(),
			ioManager,
			shuffleEnvironment,
			kvStateService,
			broadcastVariableManager,
			taskSlotTable,
			jobTable,
			jobLeaderService,
			taskStateManager,
			taskEventDispatcher,
			ioExecutor,
			libraryCacheManager
		);
	}

	private static TaskSlotTable<Task> createTaskSlotTable(
		final int numberOfSlots,
		final TaskExecutorResourceSpec taskExecutorResourceSpec,
		final long timerServiceShutdownTimeout,
		final int pageSize,
		final Executor memoryVerificationExecutor) {

		// 超时服务
		final TimerService<AllocationID> timerService = new TimerService<>(
			new ScheduledThreadPoolExecutor(1),
			timerServiceShutdownTimeout);

		// 生成 TaskSlotTableImpl
		return new TaskSlotTableImpl<>(
			// 该 TaskManager 上的 Slot 资源数
			numberOfSlots,
			// 生成资源 Profile
			TaskExecutorResourceUtils.generateTotalAvailableResourceProfile(taskExecutorResourceSpec),
			// 生成默认 Slot 资源 Profile
			TaskExecutorResourceUtils.generateDefaultSlotResourceProfile(taskExecutorResourceSpec, numberOfSlots),
			pageSize,
			timerService,
			memoryVerificationExecutor);
	}

	private static ShuffleEnvironment<?, ?> createShuffleEnvironment(TaskManagerServicesConfiguration taskManagerServicesConfiguration,
		TaskEventDispatcher taskEventDispatcher, MetricGroup taskManagerMetricGroup, Executor ioExecutor) throws FlinkException {

		// 构建 Shuffle 上下文对象：ShuffleEnvironmentContext
		final ShuffleEnvironmentContext shuffleEnvironmentContext = new ShuffleEnvironmentContext(
			taskManagerServicesConfiguration.getConfiguration(),
			taskManagerServicesConfiguration.getResourceID(),
			taskManagerServicesConfiguration.getNetworkMemorySize(),
			taskManagerServicesConfiguration.isLocalCommunicationOnly(),
			taskManagerServicesConfiguration.getBindAddress(),
			taskEventDispatcher,
			taskManagerMetricGroup,
			ioExecutor);

		// 通过反射拿到 NettyShuffleServiceFactory 对象
		return ShuffleServiceLoader.loadShuffleServiceFactory(taskManagerServicesConfiguration.getConfiguration())
			.createShuffleEnvironment(shuffleEnvironmentContext); //
	}

	/**
	 * Validates that all the directories denoted by the strings do actually exist or can be created, are proper
	 * directories (not files), and are writable.
	 *
	 * @param tmpDirs The array of directory paths to check.
	 * @throws IOException Thrown if any of the directories does not exist and cannot be created or is not writable
	 *                     or is a file, rather than a directory.
	 */
	private static void checkTempDirs(String[] tmpDirs) throws IOException {
		for(String dir : tmpDirs) {
			if(dir != null && !dir.equals("")) {
				File file = new File(dir);
				if(!file.exists()) {
					if(!file.mkdirs()) {
						throw new IOException("Temporary file directory " + file.getAbsolutePath() + " does not exist and could not be created.");
					}
				}
				if(!file.isDirectory()) {
					throw new IOException("Temporary file directory " + file.getAbsolutePath() + " is not a directory.");
				}
				if(!file.canWrite()) {
					throw new IOException("Temporary file directory " + file.getAbsolutePath() + " is not writable.");
				}

				if(LOG.isInfoEnabled()) {
					long totalSpaceGb = file.getTotalSpace() >> 30;
					long usableSpaceGb = file.getUsableSpace() >> 30;
					double usablePercentage = (double) usableSpaceGb / totalSpaceGb * 100;
					String path = file.getAbsolutePath();
					LOG.info(String
						.format("Temporary file directory '%s': total %d GB, " + "usable %d GB (%.2f%% usable)", path, totalSpaceGb, usableSpaceGb,
							usablePercentage));
				}
			} else {
				throw new IllegalArgumentException("Temporary file directory #$id is null.");
			}
		}
	}
}
